"""
增强版TimeKAN的配置类
支持新增的超参数和功能配置
"""

class EnhancedConfigs:
    """增强版TimeKAN配置类"""
    
    def __init__(self):
        # 基础配置
        self.task_name = 'long_term_forecast'
        self.is_training = 1
        self.model_id = 'TimeKAN_Enhanced'
        self.model = 'TimeKAN_Enhanced'
        
        # 数据配置
        self.data = 'ETTh1'
        self.root_path = './dataset/ETT-small/'
        self.data_path = 'ETTh1.csv'
        self.features = 'M'  # M: 多变量预测多变量, S: 单变量预测单变量, MS: 多变量预测单变量
        self.target = 'OT'
        self.freq = 'h'
        self.checkpoints = './checkpoints/'
        
        # 序列长度配置
        self.seq_len = 96
        self.label_len = 48
        self.pred_len = 96
        
        # 模型配置
        self.enc_in = 7  # 编码器输入大小
        self.dec_in = 7  # 解码器输入大小
        self.c_out = 7   # 输出大小
        self.d_model = 512  # 模型维度
        self.n_heads = 8    # 注意力头数
        self.e_layers = 2   # 编码器层数
        self.d_layers = 1   # 解码器层数
        self.d_ff = 2048    # 前馈网络维度
        self.moving_avg = 25  # 移动平均窗口
        self.factor = 1     # 注意力因子
        self.distil = True  # 是否使用蒸馏
        self.dropout = 0.1  # Dropout率
        self.embed = 'timeF'  # 时间特征编码类型
        self.activation = 'gelu'  # 激活函数
        self.output_attention = False  # 是否输出注意力
        
        # TimeKAN特有配置
        self.down_sampling_layers = 3  # 下采样层数
        self.down_sampling_window = 2  # 下采样窗口
        self.channel_independence = 1  # 通道独立性
        self.begin_order = 3  # 起始阶数
        
        # 增强功能配置
        self.use_attention = True  # 是否使用注意力机制
        self.use_multi_scale_conv = True  # 是否使用多尺度卷积
        self.use_feature_importance = True  # 是否使用特征重要性学习
        self.use_adaptive_freq = True  # 是否使用自适应频率分解
        self.use_time_aware_pos = True  # 是否使用时间感知位置编码
        
        # 注意力配置
        self.attention_dropout = 0.1  # 注意力Dropout
        self.attention_heads = 8  # 注意力头数
        
        # 多尺度卷积配置
        self.conv_kernel_sizes = [3, 5, 7]  # 卷积核大小列表
        self.conv_dropout = 0.1  # 卷积Dropout
        
        # 特征重要性配置
        self.importance_reduction_ratio = 16  # 特征重要性降维比例
        
        # 训练配置
        self.num_workers = 10
        self.itr = 1  # 实验次数
        self.train_epochs = 10  # 训练轮数
        self.batch_size = 32  # 批次大小
        self.patience = 3  # 早停耐心值
        self.learning_rate = 0.0001  # 学习率
        self.des = 'Exp'  # 实验描述
        self.loss = 'MSE'  # 损失函数
        self.lradj = 'type1'  # 学习率调整类型
        self.use_amp = False  # 是否使用混合精度
        
        # 正则化配置
        self.weight_decay = 1e-4  # 权重衰减
        self.grad_clip = 1.0  # 梯度裁剪
        
        # 其他配置
        self.use_gpu = True if torch.cuda.is_available() else False
        self.gpu = 0
        self.use_multi_gpu = False
        self.devices = '0,1,2,3'
        self.test_flop = False
        
        # 归一化配置
        self.use_norm = 1  # 是否使用归一化
        
        # 时间特征配置
        self.timeenc = 0  # 时间编码类型
        self.use_future_temporal_feature = 0  # 是否使用未来时间特征


def get_enhanced_configs(data_name='ETTh1', pred_len=96, seq_len=96):
    """
    获取增强版配置
    Args:
        data_name: 数据集名称
        pred_len: 预测长度
        seq_len: 输入序列长度
    """
    configs = EnhancedConfigs()
    
    # 根据数据集调整配置
    if data_name in ['ETTh1', 'ETTh2']:
        configs.data = data_name
        configs.data_path = f'{data_name}.csv'
        configs.enc_in = 7
        configs.dec_in = 7
        configs.c_out = 7
        configs.freq = 'h'
    elif data_name in ['ETTm1', 'ETTm2']:
        configs.data = data_name
        configs.data_path = f'{data_name}.csv'
        configs.enc_in = 7
        configs.dec_in = 7
        configs.c_out = 7
        configs.freq = 't'
    elif data_name == 'Weather':
        configs.data = 'custom'
        configs.data_path = 'weather.csv'
        configs.enc_in = 21
        configs.dec_in = 21
        configs.c_out = 21
        configs.freq = 'h'
    elif data_name == 'Electricity':
        configs.data = 'custom'
        configs.data_path = 'electricity.csv'
        configs.enc_in = 321
        configs.dec_in = 321
        configs.c_out = 321
        configs.freq = 'h'
    
    # 设置序列长度
    configs.seq_len = seq_len
    configs.pred_len = pred_len
    configs.label_len = pred_len // 2
    
    # 根据预测长度调整模型配置
    if pred_len >= 720:  # 长期预测
        configs.d_model = 768
        configs.e_layers = 3
        configs.down_sampling_layers = 4
        configs.train_epochs = 15
        configs.patience = 5
    elif pred_len >= 336:  # 中长期预测
        configs.d_model = 512
        configs.e_layers = 2
        configs.down_sampling_layers = 3
        configs.train_epochs = 12
        configs.patience = 4
    else:  # 短期预测
        configs.d_model = 256
        configs.e_layers = 2
        configs.down_sampling_layers = 2
        configs.train_epochs = 10
        configs.patience = 3
    
    return configs


def create_training_script(data_name='ETTh1', pred_len=96, seq_len=96):
    """
    创建训练脚本
    Args:
        data_name: 数据集名称
        pred_len: 预测长度
        seq_len: 输入序列长度
    """
    script_content = f"""#!/bin/bash

# 增强版TimeKAN训练脚本
# 数据集: {data_name}
# 预测长度: {pred_len}
# 输入长度: {seq_len}

export CUDA_VISIBLE_DEVICES=0

model_name=TimeKAN_Enhanced

python -u run.py \\
  --task_name long_term_forecast \\
  --is_training 1 \\
  --root_path ./dataset/ETT-small/ \\
  --data_path {data_name}.csv \\
  --model_id {data_name}_{seq_len}_{pred_len} \\
  --model $model_name \\
  --data {data_name} \\
  --features M \\
  --seq_len {seq_len} \\
  --label_len {pred_len//2} \\
  --pred_len {pred_len} \\
  --e_layers 2 \\
  --d_layers 1 \\
  --factor 3 \\
  --enc_in 7 \\
  --dec_in 7 \\
  --c_out 7 \\
  --des 'Exp' \\
  --d_model 512 \\
  --learning_rate 0.0001 \\
  --train_epochs 10 \\
  --patience 3 \\
  --down_sampling_layers 3 \\
  --down_sampling_window 2 \\
  --begin_order 3 \\
  --use_attention 1 \\
  --use_multi_scale_conv 1 \\
  --use_feature_importance 1 \\
  --use_adaptive_freq 1 \\
  --use_time_aware_pos 1 \\
  --dropout 0.1 \\
  --batch_size 32 \\
  --itr 1
"""
    
    return script_content


if __name__ == "__main__":
    import torch
    
    # 测试配置
    configs = get_enhanced_configs('ETTh1', pred_len=96, seq_len=96)
    print("增强版TimeKAN配置:")
    print(f"数据集: {configs.data}")
    print(f"模型维度: {configs.d_model}")
    print(f"预测长度: {configs.pred_len}")
    print(f"下采样层数: {configs.down_sampling_layers}")
    print(f"是否使用注意力: {configs.use_attention}")
    print(f"是否使用多尺度卷积: {configs.use_multi_scale_conv}")
    print(f"是否使用特征重要性: {configs.use_feature_importance}")
    
    # 生成训练脚本
    script = create_training_script('ETTh1', 96, 96)
    print("\n生成的训练脚本:")
    print(script[:500] + "...")